package com.example.llm.sgg.sgglangchain.config;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.example.llm.sgg.sgglangchain.po.ChatMemoryPO;
import com.example.llm.sgg.sgglangchain.service.ChatMemoryService;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.Collections;
import java.util.List;
import java.util.Objects;

/**
 * @author chenglei
 * @description: Cassandra 数据库可能更合适
 * @date 2025-05-08 12:38
 * @email 
 */
@Service
public class XiaozhiAgentChatMemoryStore implements ChatMemoryStore {

    @Autowired
    private ChatMemoryService chatMemoryService;

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        LambdaQueryWrapper<ChatMemoryPO> lambdaQueryWrapper = new LambdaQueryWrapper<>();
        lambdaQueryWrapper.eq(ChatMemoryPO::getMemoryId, memoryId.toString());
        ChatMemoryPO chatMemoryPO = chatMemoryService.getOne(lambdaQueryWrapper);
        if (Objects.isNull(chatMemoryPO)) {
            return Collections.emptyList();
        }
        return ChatMessageDeserializer.messagesFromJson(chatMemoryPO.getContent());
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        LambdaQueryWrapper<ChatMemoryPO> lambdaQueryWrapper = new LambdaQueryWrapper<>();
        lambdaQueryWrapper.eq(ChatMemoryPO::getMemoryId, memoryId.toString());
        ChatMemoryPO chatMemoryPO = chatMemoryService.getOne(lambdaQueryWrapper);
        ChatMemoryPO newChatMemoryPO;
        if (Objects.isNull(chatMemoryPO)) {
            newChatMemoryPO = new ChatMemoryPO();
            newChatMemoryPO.setMemoryId(memoryId.toString());
            newChatMemoryPO.setContent(ChatMessageSerializer.messagesToJson(messages));
        } else {
            newChatMemoryPO = chatMemoryPO;
            newChatMemoryPO.setContent(ChatMessageSerializer.messagesToJson(messages));
        }
        chatMemoryService.saveOrUpdate(newChatMemoryPO);
    }

    @Override
    public void deleteMessages(Object memoryId) {
        LambdaQueryWrapper<ChatMemoryPO> lambdaQueryWrapper = new LambdaQueryWrapper<>();
        lambdaQueryWrapper.eq(ChatMemoryPO::getMemoryId, memoryId.toString());
        chatMemoryService.remove(lambdaQueryWrapper);
    }
}
